import base64
import os
import re
import json

from Crypto.Hash import SHA1
from Crypto.PublicKey import RSA
from Crypto.Signature import pkcs1_15
from cryptography import x509
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

from step.s2_format import get_license_part_from_file

base_dir = os.path.dirname(os.path.abspath(__file__))


def generate_license_code(type="idea"):
    # type = "datagrip"
    format = get_license_part_from_file(type)
    # 正则替换日期
    format = re.sub(r"(\d{4})-(\d{2})-(\d{2})", "2099-01-01", format)
    jsonObj = json.loads(format)
    licenseId = jsonObj['licenseId']
    jsonObj['licenseeName'] = "lemon"
    licensePart = json.dumps(jsonObj)

    print(licensePart)
    # z这里是固定的加密算法
    digest = SHA1.new(licensePart.encode('utf-8'))
    # 使用第一步生成的私钥，进行加密和签名
    key_path = os.path.join(base_dir, '0/ca.key')
    # with open('0/ca.key') as prifile:
    with open(key_path) as prifile:
        private_key = RSA.import_key(prifile.read())
        # 使用私钥对HASH值进行签名
        signature = pkcs1_15.new(private_key).sign(digest)
        sig_results = base64.b64encode(signature)

        licensePartBase64 = base64.b64encode(bytes(licensePart.encode('utf-8')))
        # 此处需要使用自己的公钥进行验签
        public_key, certBase64 = get_cert_public_key()
        # 验证生成的签名是否正确
        public_key.verify(
            base64.b64decode(sig_results),
            base64.b64decode(licensePartBase64),
            padding=padding.PKCS1v15(),
            algorithm=hashes.SHA1(),
        )
        # 组装jetbrains的license格式
        result = licenseId + "-" + licensePartBase64.decode('utf-8') + "-" + sig_results.decode(
            'utf-8') + "-" + certBase64
        print(f"code: {result}")
        # 追加写入文件
        return result


def get_cert_public_key():
    ca_path = os.path.join(base_dir, '0/ca.crt')
    with open(ca_path) as prifileCrt:
        certBase64 = prifileCrt.read()
        certBase64 = certBase64.replace('-----BEGIN CERTIFICATE-----', '')
        certBase64 = certBase64.replace('-----END CERTIFICATE-----', '')
        certBase64 = certBase64.replace('\n', '')
        cert_consumer = x509.load_der_x509_certificate(base64.b64decode(certBase64))
        public_key = cert_consumer.public_key()
        return public_key, certBase64


if __name__ == '__main__':
    # clion.txt
    # datagrip.txt
    # dataspell.txt
    # dotcover.txt
    # dotmemory.txt
    # dottrace.txt
    # goland.txt
    # idea.txt
    # phpstorm.txt
    # pycharm.txt
    # rider.txt
    # rubymine.txt
    # rustrover.txt
    # webstorm.txt
    type = "idea"
    result = generate_license_code(type)
    gen_path = os.path.join(base_dir, '3/genCode.txt')
    with open(gen_path, 'a') as f:
        f.write(f"\n\n{type}")
        f.write(f"\n\n{result}\n\n")
